#include<iostream>
#include<vector>
#include<string>
#include<queue> 
#include<thread>
#include<mutex>

#include<ros/ros.h>
#include<sensor_msgs/Imu.h>
#include <geometry_msgs/Quaternion.h>
#include <tf/transform_broadcaster.h>
#include <geometry_msgs/TransformStamped.h>
#include <visualization_msgs/Marker.h>

#include<Eigen/Dense>
#include<Eigen/Geometry>
#include "sophus/so3.hpp"
#include "sophus/se2.hpp"
#include "sophus/se3.hpp"

#define IMU_HZ  200

using Vec3d = Eigen::Vector3d;
using Mat3d = Eigen::Matrix3d;

class ComplementaryFilter{
   public: 
    ComplementaryFilter( double alpha,int N):alpha_(alpha) ,N_(N) {std::cout<<"-------------INIT IMU PLEASE WAIT--------------"<<std::endl;}
    void compute(std::pair<double,std::pair<Vec3d,Vec3d>> imudata);
    void initFirstIMUPose();
    Mat3d g2R(const Vec3d &g);
    Vec3d R2ypr(const Mat3d &R);
    Mat3d ypr2R(const Vec3d &ypr);
    Vec3d GitAngle() { 
        double yaw = std::atan2(curR_(1, 0), curR_(0, 0));
        double pitch = std::atan2(-curR_(2, 0), std::sqrt(curR_(0, 0) * curR_(0, 0) + curR_(1, 0) * curR_(1, 0)));
        double roll = std::atan2(curR_(2, 1), curR_(2, 2));
        //Angle_(Vec3d(roll,pitch,yaw));
        Angle_=Vec3d(roll* 180 / M_PI,pitch* 180 / M_PI,yaw* 180 / M_PI);
        return Angle_;
    }
    Mat3d GitR()const{return curR_;}
   private:
    bool initIMUpose_ = false; 
    int N_;
    std::vector<Vec3d> accVector_;
    Vec3d averAcc_ = Vec3d(0,0,0); 
    Mat3d R0_ ,curR_; 
    std::pair<double,std::pair<Vec3d,Vec3d>> imudata_; //time gyr acc
    //Vec3d gravity(0.0, 0.0, 9.8);
    Vec3d gravity = Vec3d(0.0, 0.0, 9.8);
    double dT_ =0, pretime_ =0 , curtime_=0;
    Vec3d Angle_;//当前的角度
    Sophus::SO3d q_ , quat_ ;  //四元数角度
    double alpha_;//互补参数 0~1
};

class IMU_NODE{
    public:
        IMU_NODE(ComplementaryFilter* comple,ros::NodeHandle& nh);
        void imuDataBack(const sensor_msgs::Imu::ConstPtr& Imumsg);//ConstPtr
        std::pair<double,std::pair<Vec3d,Vec3d>> GitIMUdata() const {return IMU_data_.front();}
        void publshIMUpose(const ros::TimerEvent& event);
    private:

        ComplementaryFilter *comple_;
        ros::NodeHandle Imu_node;
        ros::Subscriber Imu_sub_;
        ros::Publisher pub_;
        ros::Timer timer_;  
        std::mutex mutex_imu_;
        //std::pair<double,std::pair<Vec3d,Vec3d>> IMU_data_; //time gyr acc
        std::queue<std::pair<double,std::pair<Vec3d,Vec3d>>> IMU_data_ ;

};
IMU_NODE::IMU_NODE(ComplementaryFilter* comple,ros::NodeHandle& nh):comple_(comple){   
    Imu_sub_ = nh.subscribe("/camera/imu",200,&IMU_NODE::imuDataBack,this);
    pub_ = nh.advertise<visualization_msgs::Marker>("IMUpose",1000);
    timer_ = nh.createTimer(ros::Duration(0.005), &IMU_NODE::publshIMUpose, this);
}
void IMU_NODE::imuDataBack(const sensor_msgs::Imu::ConstPtr& Imumsg){
    
    if(!IMU_data_.empty() &&  IMU_data_.size()>IMU_HZ){
        IMU_data_.pop();
    }

    double timestamp = Imumsg->header.stamp.toSec();
    Vec3d angular_velocity(Imumsg->angular_velocity.x,
                           Imumsg->angular_velocity.y,
                           Imumsg->angular_velocity.z);
    Vec3d linear_acceleration(Imumsg->linear_acceleration.x,
                              Imumsg->linear_acceleration.y,  //y是重力方向
                              Imumsg->linear_acceleration.z);
    mutex_imu_.lock();
    IMU_data_.push(std::make_pair(timestamp, std::make_pair(angular_velocity, linear_acceleration))); 
    mutex_imu_.unlock();
    comple_->compute(IMU_data_.front());

    // Optionally print the received IMU data
    // ROS_INFO("Received IMU data: Timestamp = %f", timestamp);
    // ROS_INFO("Angular Velocity: [%f, %f, %f]", angular_velocity[0], angular_velocity[1], angular_velocity[2]);
    // ROS_INFO("Linear Acceleration: [%f, %f, %f]", linear_acceleration[0], linear_acceleration[1], linear_acceleration[2]);  
    
}

void IMU_NODE::publshIMUpose(const ros::TimerEvent& event){
    visualization_msgs::Marker marker;
    marker.header.frame_id = "map";
    marker.header.stamp = ros::Time::now();
    marker.ns = "IMUpose";
    marker.id = 0;
    // marker.type = visualization_msgs::Marker::ARROW;  // 使用箭头表示姿态
    marker.action = visualization_msgs::Marker::ADD;
    marker.type = visualization_msgs::Marker::CUBE;

    Mat3d R = comple_->GitR();
    Eigen::Quaterniond q(R);
    q.normalize();
    geometry_msgs::Quaternion msg;
    msg.x = q.x();
    msg.y = q.y();
    msg.z = q.z();
    msg.w = q.w();

    marker.pose.orientation = msg;

    marker.pose.position.x = 0;
    marker.pose.position.y = 0;
    marker.pose.position.z = 0;
     // 设置大小
    marker.scale.x = 1.0;  
    marker.scale.y = 0.5;  
    marker.scale.z = 0.2;  

    // 设置颜色
    marker.color.r = 1.0;  // 红色
    marker.color.g = 0.0;
    marker.color.b = 0.0;
    marker.color.a = 1.0;  // 完全不透明

    pub_.publish(marker);     

}

void ComplementaryFilter::initFirstIMUPose(){
    if(accVector_.size()<N_+1){
        accVector_.push_back(imudata_.second.second);
    }
    else{
      for(size_t i = 0; i < accVector_.size();i++){
        averAcc_ += accVector_[i];
      } 
      averAcc_ = averAcc_/N_;
      std::cout<<"------------averAcc:"<<averAcc_.transpose()<<"------------"<<std::endl;
      Mat3d R0 = g2R(averAcc_);
      double yaw = R2ypr(R0).x();  
      R0 = ypr2R(Vec3d{-yaw, 0, 0}) * R0;
      curR_ = R0;//跟新为当前的姿态
      std::cout<< "R0:"<<R0<<std::endl;
      initIMUpose_ = true;
      std::cout<<"-------------INIT IMU IS END--------------"<<std::endl;

    }
    

}
Mat3d ComplementaryFilter::g2R(const Vec3d &g){
    Mat3d R0;
    Vec3d ng1 = g.normalized();
    Vec3d ng2{0, 0, 1.0};

    R0 = Eigen::Quaterniond::FromTwoVectors(ng1, ng2).toRotationMatrix();

    double yaw = ComplementaryFilter::R2ypr(R0).x();
    R0 =  ComplementaryFilter::ypr2R(Vec3d{-yaw, 0, 0}) * R0;
   
    return R0;
}
Mat3d ComplementaryFilter::ypr2R(const Vec3d &ypr){
        
        double y = ypr(0) / 180.0 * M_PI;
        double p = ypr(1) / 180.0 * M_PI;
        double r = ypr(2) / 180.0 * M_PI;

        Mat3d Rz;
        Rz << cos(y), -sin(y), 0,
            sin(y), cos(y), 0,
            0, 0, 1;

        Mat3d Ry;
        Ry << cos(p), 0., sin(p),
            0., 1., 0.,
            -sin(p), 0., cos(p);

        Mat3d Rx;
        Rx << 1., 0., 0.,
            0., cos(r), -sin(r),
            0., sin(r), cos(r);
        return Rz * Ry * Rx;
}
Vec3d ComplementaryFilter::R2ypr(const Eigen::Matrix3d &R){
        Vec3d n = R.col(0);
        Vec3d o = R.col(1);
        Vec3d a = R.col(2);
        Vec3d ypr(3);
        double y = atan2(n(1), n(0));
        double p = atan2(-n(2), n(0) * cos(y) + n(1) * sin(y));
        double r = atan2(a(0) * sin(y) - a(1) * cos(y), -o(0) * sin(y) + o(1) * cos(y));
        ypr(0) = y;
        ypr(1) = p;
        ypr(2) = r;
        return ypr / M_PI * 180.0;
}

void ComplementaryFilter::compute(std::pair<double,std::pair<Vec3d,Vec3d>> imudata){ //time gyr acc
    imudata_ = imudata;
    curtime_ = imudata_.first;
    Vec3d gyr = imudata.second.first;
    Vec3d acc = imudata.second.second;
    if(!initIMUpose_){  
        initFirstIMUPose();    
    } 
    else{
        dT_ = curtime_ - pretime_;
        Mat3d accRotation , gyrRotation;
        gyrRotation = curR_*Eigen::AngleAxisd(gyr.norm()*dT_,gyr.normalized()).toRotationMatrix() ; //AngleAxisd
        Vec3d gravityVector(0, 0, 1);
        accRotation = Eigen::Quaterniond::FromTwoVectors(acc.normalized(), gravityVector).toRotationMatrix();

        curR_ = (alpha_*gyrRotation+(1-alpha_)*accRotation).normalized();
    }
    pretime_ = curtime_;//完成时间跟新
    Vec3d angle = GitAngle();
    std::cout << "---------Angle = " <<angle.transpose()<< "------"<<std::endl;

}
double alpha = 0.15; // 0~1 需要调整

int main(int argc, char ** argv){
   
    ros::init(argc , argv , "IMU_node");
    ros::NodeHandle nh;
    ComplementaryFilter comple(alpha,10);
    IMU_NODE imu_sub_node(&comple,nh);

    ros::spin();
    
    
    return 0;
}